import asyncio
import json
import uuid
from typing import Dict, Any, Optional
from fastapi import WebSocket, WebSocketDisconnect, HTTPException, Depends
from sqlalchemy.ext.asyncio import AsyncSession

from config.get_db import get_db
from module_admin.service.login_service import LoginService
from module_speech.service.speech_service import SpeechService
from module_speech.entity.vo.speech_vo import RealtimeSpeechMessage
from utils.log_util import logger
from utils.speech_auth_util import SpeechAuthUtil


class WebSocketManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        self.active_connections: Dict[str, Dict[str, Any]] = {}
    
    async def connect(self, websocket: WebSocket, session_id: str, user_info: Dict[str, Any]):
        """建立WebSocket连接"""
        await websocket.accept()
        self.active_connections[session_id] = {
            'websocket': websocket,
            'user_info': user_info,
            'status': 'connected'
        }
        
        user_type = user_info.get('user_type', 'unknown')
        user_identifier = user_info.get('user_id') or user_info.get('client_user_id', 'unknown')
        logger.info(f"WebSocket连接建立: session_id={session_id}, user_type={user_type}, user_id={user_identifier}")
    
    def disconnect(self, session_id: str):
        """断开WebSocket连接"""
        if session_id in self.active_connections:
            del self.active_connections[session_id]
            logger.info(f"WebSocket连接断开: session_id={session_id}")
    
    async def send_message(self, session_id: str, message: Dict[str, Any]):
        """发送消息到指定会话"""
        if session_id in self.active_connections:
            websocket = self.active_connections[session_id]['websocket']
            try:
                await websocket.send_text(json.dumps(message, ensure_ascii=False))
            except Exception as e:
                logger.error(f"发送WebSocket消息失败: {str(e)}")
                self.disconnect(session_id)
    
    def get_connection_info(self, session_id: str) -> Dict[str, Any]:
        """获取连接信息"""
        return self.active_connections.get(session_id)


# 全局WebSocket管理器
ws_manager = WebSocketManager()


async def authenticate_websocket_user(token: str = None, api_key: str = None) -> Dict[str, Any]:
    """
    WebSocket用户认证
    支持JWT token（系统用户）和API Key（客户端用户）两种认证方式
    
    Args:
        token: JWT令牌
        api_key: API密钥
        
    Returns:
        Dict[str, Any]: 用户信息
        
    Raises:
        HTTPException: 认证失败时抛出异常
    """
    if token:
        # JWT认证（系统用户）
        try:
            current_user = await LoginService.verify_token(token)
            if current_user:
                return {
                    'user_type': 'system',
                    'user_id': current_user.user_id,
                    'user_info': current_user,
                    'is_admin': getattr(current_user, 'admin', False)
                }
        except Exception as e:
            logger.warning(f"JWT token验证失败: {str(e)}")
    
    if api_key:
        # API Key认证（客户端用户）
        try:
            # 创建一个模拟的request对象，因为SpeechAuthUtil.verify_api_key需要request参数
            class MockRequest:
                def __init__(self):
                    self.client = type('Client', (), {'host': '127.0.0.1'})()
                    self.headers = {'x-api-key': api_key}
            
            mock_request = MockRequest()
            client_user = await SpeechAuthUtil.verify_api_key(api_key, mock_request)
            if client_user:
                return {
                    'user_type': 'client',
                    'user_id': None,
                    'client_user_id': client_user['id'],
                    'user_info': client_user,
                    'is_admin': False
                }
        except Exception as e:
            logger.warning(f"API Key验证失败: {str(e)}")
    
    raise HTTPException(status_code=4001, detail="认证失败")


async def websocket_speech_recognize(
    websocket: WebSocket,
    session_id: str = None,
    token: str = None,
    api_key: str = None,
    db: AsyncSession = Depends(get_db)
):
    """
    WebSocket实时语音识别端点
    
    支持双用户体系认证：
    - 系统用户：使用token参数传递JWT令牌
    - 客户端用户：使用api_key参数传递API密钥
    
    连接示例：
    - 系统用户：ws://host/ws/speech?session_id=xxx&token=jwt_token
    - 客户端用户：ws://host/ws/speech?session_id=xxx&api_key=api_key
    
    消息格式:
    {
        "type": "start|data|stop|config",
        "data": {
            "model_name": "paraformer-zh-streaming",
            "language": "zh",
            "use_vad": true,
            "use_punc": true,
            "use_itn": true,
            "hotwords": ["热词1", "热词2"],
            "audio_data": "base64编码的音频数据"
        }
    }
    
    响应格式:
    {
        "type": "result|error|status",
        "session_id": "会话ID",
        "data": {
            "text": "识别结果",
            "confidence": 0.95,
            "is_final": true,
            "timestamp": "2024-01-01T12:00:00"
        }
    }
    """
    
    # 验证用户身份
    try:
        if not token and not api_key:
            await websocket.close(code=4001, reason="缺少认证参数")
            return
            
        user_info = await authenticate_websocket_user(token=token, api_key=api_key)
            
    except Exception as e:
        await websocket.close(code=4001, reason=f"身份验证失败: {str(e)}")
        return
    
    # 生成或使用现有会话ID
    if not session_id:
        session_id = str(uuid.uuid4())
    
    try:
        # 建立连接
        await ws_manager.connect(websocket, session_id, user_info)
        
        # 获取可用模型列表（根据用户类型）
        if user_info['user_type'] == 'client':
            supported_models = await SpeechService.get_available_models_for_client(
                client_user_id=user_info['client_user_id']
            )
        else:
            supported_models = await SpeechService.get_supported_models()
        
        # 发送连接成功消息
        await ws_manager.send_message(session_id, {
            "type": "status",
            "session_id": session_id,
            "data": {
                "status": "connected",
                "message": "WebSocket连接已建立",
                "user_type": user_info['user_type'],
                "supported_models": supported_models
            }
        })
        
        # 初始化识别配置
        recognition_config = {
            "model_name": "paraformer-zh-streaming",
            "language": "zh",
            "use_vad": True,
            "use_punc": True,
            "use_itn": True,
            "hotwords": []
        }
        
        # 消息处理循环
        while True:
            try:
                # 接收客户端消息
                message_text = await websocket.receive_text()
                message = json.loads(message_text)
                
                await handle_websocket_message(
                    ws_manager,
                    session_id,
                    user_info,
                    message,
                    recognition_config,
                    db
                )
                
            except WebSocketDisconnect:
                logger.info(f"WebSocket客户端主动断开连接: session_id={session_id}")
                break
                
            except json.JSONDecodeError as e:
                await ws_manager.send_message(session_id, {
                    "type": "error",
                    "session_id": session_id,
                    "data": {
                        "error": "消息格式错误",
                        "detail": str(e)
                    }
                })
                
            except Exception as e:
                logger.error(f"处理WebSocket消息失败: {str(e)}")
                await ws_manager.send_message(session_id, {
                    "type": "error",
                    "session_id": session_id,
                    "data": {
                        "error": "处理消息失败",
                        "detail": str(e)
                    }
                })
    
    except Exception as e:
        logger.error(f"WebSocket连接异常: {str(e)}")
        
    finally:
        # 清理连接和会话
        ws_manager.disconnect(session_id)
        await SpeechService.clear_session(session_id)


async def handle_websocket_message(
    ws_manager: WebSocketManager,
    session_id: str,
    user_info: Dict[str, Any],
    message: Dict[str, Any],
    recognition_config: Dict[str, Any],
    db: AsyncSession
):
    """处理WebSocket消息"""
    
    message_type = message.get("type")
    data = message.get("data", {})
    
    if message_type == "config":
        # 更新识别配置
        await handle_config_message(ws_manager, session_id, user_info, data, recognition_config)
        
    elif message_type == "start":
        # 开始识别
        await handle_start_message(ws_manager, session_id, user_info, data, recognition_config, db)
        
    elif message_type == "data":
        # 音频数据
        await handle_data_message(ws_manager, session_id, user_info, data, recognition_config)
        
    elif message_type == "stop":
        # 停止识别
        await handle_stop_message(ws_manager, session_id, user_info, data, recognition_config)
        
    else:
        await ws_manager.send_message(session_id, {
            "type": "error",
            "session_id": session_id,
            "data": {
                "error": "未知的消息类型",
                "message_type": message_type
            }
        })


async def handle_config_message(
    ws_manager: WebSocketManager,
    session_id: str,
    user_info: Dict[str, Any],
    data: Dict[str, Any],
    recognition_config: Dict[str, Any]
):
    """处理配置更新消息"""
    try:
        # 验证并更新配置
        model_name = data.get("model_name")
        if model_name:
            # 检查用户是否有权限使用指定模型
            if user_info['user_type'] == 'client':
                available_models = await SpeechService.get_available_models_for_client(
                    client_user_id=user_info['client_user_id']
                )
                if model_name not in available_models:
                    await ws_manager.send_message(session_id, {
                        "type": "error",
                        "session_id": session_id,
                        "data": {
                            "error": "无权限使用指定模型",
                            "model_name": model_name,
                            "available_models": available_models
                        }
                    })
                    return
            
            recognition_config["model_name"] = model_name
        
        # 更新其他配置
        for key in ["language", "use_vad", "use_punc", "use_itn", "hotwords"]:
            if key in data:
                recognition_config[key] = data[key]
        
        await ws_manager.send_message(session_id, {
            "type": "status",
            "session_id": session_id,
            "data": {
                "status": "config_updated",
                "config": recognition_config
            }
        })
        
    except Exception as e:
        await ws_manager.send_message(session_id, {
            "type": "error",
            "session_id": session_id,
            "data": {
                "error": "配置更新失败",
                "detail": str(e)
            }
        })


async def handle_start_message(
    ws_manager: WebSocketManager,
    session_id: str,
    user_info: Dict[str, Any],
    data: Dict[str, Any],
    recognition_config: Dict[str, Any],
    db: AsyncSession
):
    """处理开始识别消息"""
    try:
        # 检查用户配额（仅客户端用户）
        if user_info['user_type'] == 'client':
            from module_speech.service.speech_quota_service import QuotaService
            quota_available = await QuotaService.check_monthly_quota(
                db=db,
                client_user_id=user_info['client_user_id']
            )
            if not quota_available:
                await ws_manager.send_message(session_id, {
                    "type": "error",
                    "session_id": session_id,
                    "data": {
                        "error": "月度配额不足",
                        "error_code": "QUOTA_EXCEEDED"
                    }
                })
                return
        
        # 初始化实时识别
        success = await SpeechService.start_realtime_recognition(
            session_id=session_id,
            config=recognition_config,
            user_id=user_info.get('user_id'),
            client_user_id=user_info.get('client_user_id'),
            user_type=user_info['user_type']
        )
        
        if success:
            await ws_manager.send_message(session_id, {
                "type": "status",
                "session_id": session_id,
                "data": {
                    "status": "recognition_started",
                    "config": recognition_config
                }
            })
        else:
            await ws_manager.send_message(session_id, {
                "type": "error",
                "session_id": session_id,
                "data": {
                    "error": "启动识别失败"
                }
            })
            
    except Exception as e:
        await ws_manager.send_message(session_id, {
            "type": "error",
            "session_id": session_id,
            "data": {
                "error": "启动识别异常",
                "detail": str(e)
            }
        })


async def handle_data_message(
    ws_manager: WebSocketManager,
    session_id: str,
    user_info: Dict[str, Any],
    data: Dict[str, Any],
    recognition_config: Dict[str, Any]
):
    """处理音频数据消息"""
    try:
        audio_data = data.get("audio_data")
        if not audio_data:
            await ws_manager.send_message(session_id, {
                "type": "error",
                "session_id": session_id,
                "data": {
                    "error": "缺少音频数据"
                }
            })
            return
        
        # 处理实时音频数据
        result = await SpeechService.process_realtime_audio_data(
            session_id=session_id,
            audio_data=audio_data,
            user_id=user_info.get('user_id'),
            client_user_id=user_info.get('client_user_id'),
            user_type=user_info['user_type']
        )
        
        if result:
            await ws_manager.send_message(session_id, {
                "type": "result",
                "session_id": session_id,
                "data": result
            })
            
    except Exception as e:
        await ws_manager.send_message(session_id, {
            "type": "error",
            "session_id": session_id,
            "data": {
                "error": "处理音频数据失败",
                "detail": str(e)
            }
        })


async def handle_stop_message(
    ws_manager: WebSocketManager,
    session_id: str,
    user_info: Dict[str, Any],
    data: Dict[str, Any],
    recognition_config: Dict[str, Any]
):
    """处理停止识别消息"""
    try:
        # 停止实时识别并获取最终结果
        final_result = await SpeechService.stop_realtime_recognition(
            session_id=session_id,
            user_id=user_info.get('user_id'),
            client_user_id=user_info.get('client_user_id'),
            user_type=user_info['user_type']
        )
        
        await ws_manager.send_message(session_id, {
            "type": "result",
            "session_id": session_id,
            "data": {
                "status": "final",
                "result": final_result,
                "session_completed": True
            }
        })
        
    except Exception as e:
        await ws_manager.send_message(session_id, {
            "type": "error",
            "session_id": session_id,
            "data": {
                "error": "停止识别失败",
                "detail": str(e)
            }
        })


def get_websocket_manager() -> WebSocketManager:
    """获取WebSocket管理器实例"""
    return ws_manager 